import asyncio
import os
from time import sleep

import cv2
import matplotlib.pyplot
from fastapi import UploadFile

from Imgprocesser import ImageProcessor
from ultralytics import YOLO


class YOLOProcessor:
    def __init__(self, model_path):
        self.model_path = model_path
        self.model = YOLO(self.model_path)
        with open('coco.names', 'r') as f:
            self.class_names = eval(f.read())

    async def process_with_yolo(self, img: UploadFile):
        file_location = f"../temp/{img.filename}"
        with open(file_location, "wb+") as file_object:
            file_object.write(img.file.read())

        await asyncio.sleep(1)  # 使用异步sleep
        result = self.model.predict(file_location, save=True)
        print(result)
        max_confidence = 0
        max_class_id = None
        for r in result:
            boxes = r.boxes  # Boxes object for bbox outputs
            class_ids = boxes.cls.cpu().numpy().astype(int)  # 转为int类型数组
            confidences = boxes.conf.cpu().numpy()  # 转为numpy数组
            # 找到置信度最高的类别ID
            if confidences.size > 0 and max(confidences) > max_confidence:
                max_confidence = max(confidences)
                max_class_id = class_ids[confidences.argmax()]
        print(f"最可能的类别ID是：{max_class_id}")
        class_name = self.class_names[max_class_id] if max_class_id is not None else None
        print(f"最可能的类别名称是：{class_name}")
        img_list = self.list_images()
        if img_list:
            img_path = f'./runs/detect/predict/{img_list[-1]}'
            print(img_path)
            base64 = ImageProcessor.image_to_byte_array(img_path)
            return {"data": base64, "id": int(max_class_id), "name": class_name}
        else:
            return None

    @staticmethod
    def list_images():
        folder_path = './runs/detect/predict'
        image_extensions = ['.jpg', '.jpeg', '.png']
        return [file for file in os.listdir(folder_path) if os.path.splitext(file)[1].lower() in image_extensions]
